package com.youngjun.common.util;

import org.apache.commons.io.IOUtils;
import org.bouncycastle.pkcs.PKCS10CertificationRequest;
import java.io.*;
import java.nio.charset.Charset;
import java.security.KeyFactory;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.cert.CertificateEncodingException;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

public class PKIUtils {

    private final static Charset DEFAULT_Charset = Charset.forName("UTF-8");

    private final static String[] PRIVATE_KEY_FLAG = {
            "-----BEGIN RSA PRIVATE KEY-----",
            "-----END RSA PRIVATE KEY-----",
            "-----BEGIN PRIVATE KEY-----",
            "-----END PRIVATE KEY-----",
            "\\t",
            "\\n",
            "\t",
            "\n",
    };
    private final static String[] PUBLIC_KEY_FLAG = {
            "-----BEGIN PUBLIC KEY-----",
            "-----END PUBLIC KEY-----",
            "\\t",
            "\\n",
            "\t",
            "\n",
    };

    private final static String[] CERT_FLAG = {
            "-----BEGIN CERTIFICATE-----",
            "-----END CERTIFICATE-----",
            "\\t",
            "\\n",
            "\t",
            "\n",
    };
    private final static String[] CERT_REQUEST = {
            "-----BEGIN CERTIFICATE REQUEST-----",
            "-----END CERTIFICATE REQUEST-----",
            "\\t",
            "\\n",
            "\t",
            "\n",
    };

    public static PKCS10CertificationRequest loadCertificateSigningRequest(String pemBase64EncodeStr) throws IOException {
        return new PKCS10CertificationRequest(parseFromPEM(CERT_REQUEST, pemBase64EncodeStr));
    }

    public static X509Certificate loadCertificateFromDER(String type, InputStream inputStream) throws CertificateException {
        return (X509Certificate) CertificateFactory.getInstance(type).generateCertificate(inputStream);
    }

    public static X509Certificate loadCertificateFromPEM(String type, InputStream inputStream) throws CertificateException, IOException {
        String certBase64Str = IOUtils.toString(inputStream);
        byte[] bytes = parseFromPEM(CERT_FLAG, certBase64Str);
        return (X509Certificate) CertificateFactory.getInstance(type).generateCertificate(new ByteArrayInputStream(bytes));
    }

    public static PrivateKey loadPrivateKeyFromDER(byte[] derBytes, String algorithm) throws NoSuchAlgorithmException, InvalidKeySpecException {
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(derBytes);
        KeyFactory kf = KeyFactory.getInstance(algorithm);
        return kf.generatePrivate(keySpec);
    }

    public static PrivateKey loadPrivateKeyFromPEM(String pemBase64EncodeStr, String algorithm) throws NoSuchAlgorithmException, InvalidKeySpecException {
        byte[] bytes = parseFromPEM(PRIVATE_KEY_FLAG, pemBase64EncodeStr);
        return loadPrivateKeyFromDER(bytes, algorithm);
    }

    public static PublicKey loadPublicKeyFromDER(byte[] derBytes, String algorithm) throws NoSuchAlgorithmException, InvalidKeySpecException {
        return KeyFactory.getInstance(algorithm).generatePublic(new X509EncodedKeySpec(derBytes));
    }

    public static PublicKey loadPublicKeyFromPEM(String pemBase64EncodeStr, String algorithm) throws InvalidKeySpecException, NoSuchAlgorithmException {
        byte[] bytes = parseFromPEM(PUBLIC_KEY_FLAG, pemBase64EncodeStr);
        return loadPublicKeyFromDER(bytes, algorithm);
    }

    public static String exportToPEM(PublicKey publicKey) throws IOException {
        return convertFormatToPEM(PUBLIC_KEY_FLAG, publicKey.getEncoded());
    }

    public static byte[] exportToDER(PublicKey publicKey) {
        return publicKey.getEncoded();
    }

    public static String exportToPEM(PrivateKey privateKey) throws IOException {
        return convertFormatToPEM(PRIVATE_KEY_FLAG, privateKey.getEncoded());
    }

    public static byte[] exportToDER(PrivateKey privateKey) {
        return privateKey.getEncoded();
    }

    public static String exportToPEM(X509Certificate certificate) throws IOException, CertificateEncodingException {
        return convertFormatToPEM(CERT_FLAG, certificate.getEncoded());
    }


    public static byte[] exportToDER(X509Certificate certificate) throws CertificateEncodingException {
        return certificate.getEncoded();
    }


    private static String convertFormatToPEM(String[] flags, byte[] encoded) throws IOException {
        try (
                StringWriter stringWriter = new StringWriter();
                BufferedWriter out = new BufferedWriter(stringWriter);
        ) {
            out.write(flags[0]); //write header to the file
            out.newLine(); //new line
            //write cert data to the file, after Base64 encoding it
            out.write(Base64.getMimeEncoder().encodeToString(encoded));
            out.newLine(); //new line
            out.write(flags[1]); //write footer
            out.flush();
            return stringWriter.toString();
        }
    }

    private static byte[] parseFromPEM(String[] flags, String pemBase64EncodeStr) {
        for (String flag : flags) {
            pemBase64EncodeStr = pemBase64EncodeStr.replace(flag, "");
        }
        return Base64.getMimeDecoder().decode(pemBase64EncodeStr);
    }
}
